# -*- coding: utf-8 -*-
"""
完整可运行的递归分块 Demo
包含：分块、向量化、检索、问答完整流程
"""

import os
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import textwrap


# =======================
# 1. 准备示例数据
# =======================
def create_sample_data():
    """创建示例文档数据"""

    documents = [
        Document(
            page_content="""
人工智能技术的发展历程

人工智能的概念最早可以追溯到1950年代。1956年，在达特茅斯会议上，约翰·麦卡锡首次提出了"人工智能"这一术语，标志着AI领域的正式诞生。

早期AI研究主要集中于符号推理和专家系统。研究人员试图通过编写规则来模拟人类智能，但在处理不确定性和复杂模式识别方面遇到了困难。

1980年代，机器学习开始兴起。与基于规则的方法不同，机器学习让计算机从数据中自动学习模式。这一时期的代表性算法包括决策树、支持向量机等。

2010年代，深度学习革命彻底改变了AI领域。由于大数据、强大算力（GPU）和改进算法的出现，神经网络在图像识别、自然语言处理等任务上取得了突破性进展。

近年来，大语言模型（如GPT系列）和生成式AI技术快速发展，使AI在创造性任务、对话交互等方面展现出惊人能力。
""",
            metadata={"source": "ai_history", "category": "技术历史"}
        ),

        Document(
            page_content="""
机器学习基础概念

机器学习是人工智能的一个分支，它使计算机能够在没有明确编程的情况下学习。机器学习算法可以分为三大类：

监督学习：使用带标签的数据训练模型，用于分类和回归任务。常见算法包括线性回归、逻辑回归、支持向量机等。

无监督学习：处理没有标签的数据，用于发现数据中的内在模式。常见算法包括K均值聚类、主成分分析等。

强化学习：智能体通过与环境交互学习最优策略，以最大化累积奖励。这种方法在游戏AI和机器人控制中很受欢迎。

深度学习是机器学习的一个子领域，它使用多层神经网络来学习数据的层次化表示。卷积神经网络（CNN）擅长处理图像数据，循环神经网络（RNN）适合序列数据，Transformer架构则在自然语言处理中占主导地位。
""",
            metadata={"source": "ml_basics", "category": "基础知识"}
        ),

        Document(
            page_content="""
大语言模型的工作原理

大语言模型（LLM）是基于Transformer架构的深度学习模型，通过在大量文本数据上进行预训练来学习语言的统计规律。

Transformer架构的核心是自注意力机制，它允许模型在处理每个词时考虑输入序列中的所有其他词，从而更好地理解上下文关系。

预训练过程通常使用掩码语言建模目标，即随机遮盖输入文本中的一些词，让模型预测被遮盖的词。通过这个过程，模型学习到语法、语义和世界知识。

微调是在预训练基础上，使用特定任务的数据对模型进行进一步训练，使其适应具体应用场景。指令微调让模型更好地理解并遵循人类指令。

RAG（检索增强生成）技术将大语言模型与外部知识库结合，在生成回答前先检索相关信息，从而提高回答的准确性和时效性。
""",
            metadata={"source": "llm_tech", "category": "前沿技术"}
        ),

        Document(
            page_content="""
实际应用场景

在客服领域，AI助手可以处理常见问题查询，提供24/7服务，大大降低人力成本。智能客服能够理解用户意图，提供准确的解决方案。

在教育行业，AI可以个性化学习路径，根据学生的学习进度和理解程度调整教学内容。智能辅导系统能够即时回答学生问题，提供详细解释。

在医疗健康领域，AI辅助诊断系统可以帮助医生分析医学影像，识别病变特征。药物研发中，AI可以加速新药发现过程，预测分子性质。

在创意产业，生成式AI可以协助创作文本、图像、音乐等内容。写作助手、设计工具和音乐生成平台正在改变创意工作流程。

企业应用包括智能文档处理、数据分析、预测维护等。AI技术正在各个行业提高效率，创造新的商业价值。
""",
            metadata={"source": "applications", "category": "应用实践"}
        )
    ]

    return documents


# =======================
# 2. 递归分块处理
# =======================
def setup_recursive_chunking(documents, chunk_size=500, chunk_overlap=50):
    """设置并执行递归分块"""

    print("🔧 正在设置递归分块器...")

    # 创建递归分块器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""]
    )

    # 对文档进行分块
    all_chunks = []
    for doc in documents:
        chunks = text_splitter.split_text(doc.page_content)
        for chunk in chunks:
            # 为每个分块创建新文档，保留元数据
            new_doc = Document(page_content=chunk, metadata=doc.metadata.copy())
            all_chunks.append(new_doc)

    print(f"✅ 分块完成: {len(documents)} 篇文档 → {len(all_chunks)} 个分块")

    # 显示分块统计信息
    chunk_lengths = [len(chunk.page_content) for chunk in all_chunks]
    avg_length = sum(chunk_lengths) / len(chunk_lengths)

    print(f"📊 分块统计:")
    print(f"   - 平均长度: {avg_length:.1f} 字符")
    print(f"   - 最小长度: {min(chunk_lengths)} 字符")
    print(f"   - 最大长度: {max(chunk_lengths)} 字符")
    print(f"   - 配置: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}")

    return all_chunks


# =======================
# 3. 创建向量数据库
# =======================
def create_vector_store(chunks):
    """创建向量数据库"""

    print("\n🔧 正在初始化嵌入模型...")

    # 使用本地嵌入模型（如果不可用，会回退到其他模型）

    embeddings = HuggingFaceEmbeddings(model_name="D:/models/BAAIbge-base-zh-v1.5", model_kwargs={'device': 'cpu'})
    print("🔧 正在创建向量数据库...")

    # 创建 FAISS 向量数据库
    vector_store = FAISS.from_documents(chunks, embeddings)

    print("✅ 向量数据库创建完成")
    print(f"📊 索引中包含 {vector_store.index.ntotal} 个向量")

    return vector_store


# =======================
# 4. 设置问答系统
# =======================
def setup_qa_system(vector_store):
    """设置问答系统"""

    print("\n🔧 正在初始化语言模型...")

    # 初始化语言模型
    # llm = ChatOpenAI(model="deepseek-ai/DeepSeek-R1-0528", temperature=0.7)
    llm = ChatOpenAI(model="xdeepseekv32exp", temperature=0.7)

    # 创建自定义提示模板
    prompt_template = PromptTemplate(
        template=textwrap.dedent("""
            请根据以下上下文信息回答问题。如果上下文中没有足够信息，请如实告知。

            上下文：
            {context}

            问题：{question}

            请提供准确、详细的回答：
            """),
        input_variables=["context", "question"]
    )

    # 创建检索式QA链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
        chain_type_kwargs={"prompt": prompt_template},
        return_source_documents=True
    )

    print("✅ 问答系统设置完成")
    return qa_chain


# =======================
# 5. 测试函数
# =======================
def test_retrieval(vector_store):
    """测试检索功能"""

    print("\n🧪 测试检索功能...")

    test_queries = [
        "什么是机器学习？",
        "大语言模型如何工作？",
        "AI在医疗领域有哪些应用？",
        "Transformer架构的特点是什么？"
    ]

    retriever = vector_store.as_retriever(search_kwargs={"k": 2})

    for query in test_queries:
        print(f"\n🔍 查询: '{query}'")
        docs = retriever.invoke(query)

        for i, doc in enumerate(docs):
            source = doc.metadata.get('source', '未知')
            content_preview = doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content
            print(f"   结果 {i + 1} (来源: {source}): {content_preview}")


# =======================
# 6. 交互式问答
# =======================
def interactive_qa(qa_chain):
    """交互式问答循环"""

    print("\n" + "=" * 60)
    print("🤖 问答系统已就绪！输入您的问题（输入 'quit' 退出）")
    print("=" * 60)

    while True:
        question = input("\n❓ 请输入问题: ").strip()

        if question.lower() in ['quit', 'exit', '退出']:
            print("👋 再见！")
            break

        if not question:
            continue

        try:
            print("🔍 正在检索和生成回答...")
            result = qa_chain.invoke({"query": question})

            print("\n💡 回答:")
            print("-" * 50)
            print(result["result"])

            # 显示来源文档
            print("\n📚 参考来源:")
            for i, doc in enumerate(result["source_documents"]):
                source = doc.metadata.get('source', '未知')
                category = doc.metadata.get('category', '未知')
                print(f"  {i + 1}. [{category}] {source}")
                # 显示文档片段预览
                content_preview = textwrap.shorten(doc.page_content, width=80, placeholder="...")
                print(f"     片段: {content_preview}")
            print("-" * 50)

        except Exception as e:
            print(f"❌ 出错: {e}")


# =======================
# 7. 主函数
# =======================
def main():
    """主函数"""

    print("🚀 开始递归分块 RAG 系统演示")
    print("=" * 60)

    # 步骤1: 准备数据
    print("📝 步骤 1/4: 准备示例数据")
    documents = create_sample_data()
    print(f"✅ 已创建 {len(documents)} 篇示例文档")

    # 步骤2: 递归分块
    print("\n📝 步骤 2/4: 执行递归分块")
    chunks = setup_recursive_chunking(documents, chunk_size=400, chunk_overlap=50)

    # 步骤3: 创建向量数据库
    print("\n📝 步骤 3/4: 创建向量数据库")
    vector_store = create_vector_store(chunks)

    # 步骤4: 设置问答系统
    print("\n📝 步骤 4/4: 设置问答系统")
    qa_chain = setup_qa_system(vector_store)

    # 测试检索
    test_retrieval(vector_store)

    # 开始交互式问答
    interactive_qa(qa_chain)


# =======================
# 8. 额外工具函数
# =======================
def analyze_chunking_effectiveness():
    """分析不同分块参数的效果"""

    print("\n📊 分块参数效果分析")
    print("=" * 50)

    documents = create_sample_data()

    # 测试不同参数
    parameters = [
        (300, 30, "小分块"),
        (500, 50, "中等分块"),
        (800, 80, "大分块")
    ]

    for chunk_size, chunk_overlap, desc in parameters:
        chunks = setup_recursive_chunking(
            documents,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap
        )

        # 简单分析分块质量
        lengths = [len(chunk.page_content) for chunk in chunks]
        uniformity = max(lengths) - min(lengths)  # 长度差异越小越均匀

        print(f"  {desc}: {len(chunks)} 块, 长度差异: {uniformity}")


if __name__ == "__main__":
    # 运行主演示
    main()

    # 可选：运行参数分析
    # analyze_chunking_effectiveness()
